import socket
import sys
import random
import time
import threading
import pandas as pd
from datetime import datetime


class ReliableUDPClient:
    def __init__(self, server_ip, server_port):
        self.server_address = (server_ip, server_port)
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.client_socket.settimeout(0.3)  # 超时时间300ms

        # 要发送的长数据
        self.data_content = "I am the data from client to server" * 20 + "END"
        self.total_data_size = len(self.data_content)
        self.data_index = 0

        # GBN协议参数
        self.window_size = 400  # 发送窗口大小400字节
        self.packet_size_min = 40  # 数据包最小大小
        self.packet_size_max = 80  # 数据包最大大小
        self.max_packets = 30  # 总数据包数

        # 窗口控制
        self.base_seq = 1  # 窗口起始序号
        self.next_seq = 1  # 下一个要发送的序号
        self.timeout_threshold = 0.3  # 初始超时阈值
        self.last_ack = 0  # 最后收到的ACK号

        #RTT的计算
        self.srtt = None  # 平滑RTT (毫秒)
        self.rttvar = None  # RTT变化量 (毫秒)
        self.rto = 1.0  # 初始重传超时 (秒)
        # 初始化参数
        self.ALPHA = 0.125  # SRTT平滑因子
        self.BETA = 0.25  # RTTVAR平滑因子
        self.K = 4  # RTO乘数因子
        self.G = 1.0  # 最小粒度(ms)

        # 数据管理
        self.send_buffer = {}  # 发送缓冲区
        self.rtt_samples = []  # RTT样本
        self.packets_sent = 0  # 发送的数据包总数
        self.packets_retransmitted = 0  # 重传的数据包数
        self.start_time = None  # 开始时间
        self.connected = False  # 连接状态
        self.lock = threading.Lock()  # 用于线程安全的锁
        self.receiver_thread = None
        self.sender_thread = None
        self.fin_event = threading.Event()  # 用于通知关闭
        self.ack_received = threading.Event()  # 用于通知ACK到达
        self.dup_ack_count = {}  # 重复ACK计数器

    def connect(self):
        """实现三次握手建立连接"""
        # 生成客户端初始序列号
        client_isn = random.randint(1000, 5000)
        syn_packet = self.create_packet(
            seq_num=client_isn,
            flags='SYN'
        )

        retries = 0
        while retries < 3 and not self.connected:
            try:
                self.client_socket.sendto(syn_packet, self.server_address)
                print(f"[CLIENT] 发送SYN，seq={client_isn}，等待服务器响应...")

                # 接收服务器的SYN+ACK
                data, addr = self.client_socket.recvfrom(1024)
                if addr != self.server_address:
                    continue
                seq_num, ack_num, flags, _ = self.parse_packet(data)

                if flags == 'SYN+ACK' and ack_num == client_isn + 1:
                    print(f"[CLIENT] 收到SYN+ACK，seq={seq_num}, ack={ack_num}")

                    # 发送ACK完成三次握手
                    ack_packet = self.create_packet(
                        seq_num=client_isn + 1,
                        ack_num=seq_num + 1,
                        flags='ACK'
                    )
                    self.client_socket.sendto(ack_packet, self.server_address)
                    print(f"[CLIENT] 发送ACK，完成三次握手，seq={client_isn + 1}, ack={seq_num + 1}")

                    # 等待服务器确认连接
                    data, addr = self.client_socket.recvfrom(1024)
                    if addr != self.server_address:
                        continue
                    seq_num, ack_num, flags, _ = self.parse_packet(data)
                    if flags == 'CONNECTED':
                        print("[CLIENT] 连接已建立！")
                        self.connected = True
                        return True
            except socket.timeout:
                retries += 1
                print(f"[CLIENT] SYN超时，重试 {retries}/3...")
        print("[CLIENT] 连接建立失败")
        return False

    def start_transfer(self):
        """启动数据传输"""
        if not self.connected:
            print("[CLIENT] 错误：未建立连接，无法发送数据")
            return

        self.start_time = time.time()
        print("[CLIENT] 开始发送数据...")
        print(f"[CLIENT] 总数据大小: {self.total_data_size}字节，将分成最多{self.max_packets}个包")

        # 启动接收线程
        self.receiver_thread = threading.Thread(target=self.receive_acks)
        self.receiver_thread.daemon = True
        self.receiver_thread.start()

        # 启动发送线程
        self.sender_thread = threading.Thread(target=self.send_data)
        self.sender_thread.start()

        # 等待发送线程完成
        self.sender_thread.join()  # 阻塞直到发送线程结束

        # 设置结束事件
        self.fin_event.set()

        # 等待接收线程结束
        if self.receiver_thread:
            self.receiver_thread.join(timeout=2.0)

        # 等待接收线程处理完所有ACK
        time.sleep(1)

        # 发送FIN关闭连接
        self.close()

        # 打印统计信息
        self.print_summary()

    def send_data(self):
        """发送数据线程"""

        # 计算当前窗口内已使用的字节数
        def window_used():
            with self.lock:
                used = 0
                for seq in range(self.base_seq, self.next_seq):
                    if seq in self.send_buffer:
                        pkt = self.send_buffer[seq]
                        used += pkt['size']
                return used

        # 发送所有数据包
        while self.data_index < self.total_data_size and self.next_seq <= self.max_packets:
            if self.fin_event.is_set():
                break

            # 检查窗口是否有空间
            current_used = window_used()
            if current_used < self.window_size:
                # 计算数据包大小
                remaining_size = self.total_data_size - self.data_index
                packet_size = min(
                    random.randint(self.packet_size_min, self.packet_size_max),
                    remaining_size
                )

                # 确保不超过窗口大小
                if current_used + packet_size > self.window_size:
                    # 如果添加后会超过窗口大小，则等待ACK
                    self.ack_received.wait(0.5)
                    self.ack_received.clear()
                    continue

                # 记录数据边界
                start_byte = self.data_index
                end_byte = self.data_index + packet_size - 1

                # 获取数据分段
                data_segment = self.data_content[self.data_index:self.data_index + packet_size]
                self.data_index += packet_size

                # 创建数据包
                packet = self.create_packet(
                    seq_num=self.next_seq,
                    flags='DATA',
                    data=data_segment
                )

                # 记录发送信息
                with self.lock:
                    self.send_buffer[self.next_seq] = {
                        'packet': packet,
                        'sent_time': time.time(),
                        'retries': 0,
                        'acked': False,
                        'start_byte': start_byte,
                        'end_byte': end_byte,
                        'size': packet_size
                    }

                # 发送数据包
                self.client_socket.sendto(packet, self.server_address)
                self.packets_sent += 1

                # 打印发送信息
                print(f"[CLIENT] 第{self.next_seq}个（第{start_byte}~{end_byte}字节）client端已经发送")
                print(f"[CLIENT] 发送数据: {data_segment[:20]}{'...' if len(data_segment) > 20 else ''}")

                # 添加发送延迟 (1-5ms随机)
                time.sleep(random.uniform(0.001, 0.005))

                self.next_seq += 1
            else:
                # 窗口已满，等待ACK
                self.ack_received.wait(0.5)
                self.ack_received.clear()

        # 等待所有数据包确认
        while self.base_seq < self.next_seq and not self.fin_event.is_set():
            time.sleep(0.1)

    def receive_acks(self):
        """接收ACK的线程"""
        while not self.fin_event.is_set():
            try:
                # 设置较短的超时 (100ms)
                self.client_socket.settimeout(0.1)

                data, addr = self.client_socket.recvfrom(1024)
                if addr != self.server_address:
                    continue

                seq_num, ack_num, flags, _ = self.parse_packet(data)

                if flags == 'ACK':
                    # 处理ACK和重复ACK
                    self.handle_ack(ack_num)

                    # 通知发送线程有ACK到达
                    self.ack_received.set()

                    # 添加处理延迟 (0.5-1ms)
                    time.sleep(random.uniform(0.0005, 0.001))
                elif flags == 'FIN':
                    self.handle_server_fin(seq_num)
            except socket.timeout:
                # 检查是否有超时需要重传
                self.retransmit_timeout()
            except Exception as e:
                if not self.fin_event.is_set():
                    print(f"[CLIENT] 接收ACK时出错: {e}")
                break

    def handle_ack(self, ack_num):
        """处理ACK - 支持快重传"""
        with self.lock:
            if ack_num > self.base_seq:
                # 处理新的ACK
                self.process_new_ack(ack_num)
            elif ack_num == self.base_seq and ack_num > 0:
                # 处理重复ACK
                self.process_dup_ack(ack_num)

    def process_new_ack(self, ack_num):
        """处理新的ACK"""
        # 重置重复ACK计数器
        self.dup_ack_count = {}

        # 计算RTT并更新确认状态
        for seq in range(self.base_seq, ack_num):
            if seq in self.send_buffer and not self.send_buffer[seq]['acked']:
                rtt = (time.time() - self.send_buffer[seq]['sent_time']) * 1000
                self.rtt_samples.append(rtt)
                # 调整超时时间
                self.adjust_timeout(rtt)

                self.send_buffer[seq]['acked'] = True

                # 打印确认信息
                server_time = datetime.now().strftime("%H:%M:%S")
                print(f"[CLIENT] 第{seq}个包已确认, RTT={rtt:.2f}ms, 时间:{server_time}")

        # 滑动窗口
        old_base = self.base_seq
        self.base_seq = ack_num
        print(f"[CLIENT] 窗口由 {old_base} 滑动到 {self.base_seq}")

        # # 调整超时时间
        # self.adjust_timeout()

        # 更新最后收到的ACK
        self.last_ack = ack_num

    def process_dup_ack(self, ack_num):
        """处理重复ACK - 实现快重传"""
        # 更新重复ACK计数
        if ack_num in self.dup_ack_count:
            self.dup_ack_count[ack_num] += 1
        else:
            self.dup_ack_count[ack_num] = 1

        print(f"[CLIENT] 收到重复ACK {ack_num} (计数: {self.dup_ack_count[ack_num]})")

        # 快重传机制 (收到3个重复ACK时触发)
        if self.dup_ack_count[ack_num] >= 3:
            print(f"[CLIENT] 触发快重传! 重传第{ack_num}个包")
            self._retransmit_packet(ack_num)

            # 重置计数器
            self.dup_ack_count[ack_num] = 0

    def _retransmit_packet(self, seq):
        """重传单个数据包"""
        if seq in self.send_buffer:
            pkt = self.send_buffer[seq]
            self.client_socket.sendto(pkt['packet'], self.server_address)
            pkt['sent_time'] = time.time()
            pkt['retries'] += 1
            self.packets_sent += 1
            self.packets_retransmitted += 1
            print(f"[CLIENT] 重传第{seq}个数据包")

            # 添加重传延迟 (2-5ms)
            time.sleep(random.uniform(0.002, 0.005))

    def retransmit_timeout(self):
        """处理超时重传 - 优化版"""
        with self.lock:
            current_time = time.time()
            timeout_occurred = False

            # 只检查最早发送的未确认包
            earliest_unacked = None
            for seq in range(self.base_seq, self.next_seq):
                if seq in self.send_buffer and not self.send_buffer[seq]['acked']:
                    pkt = self.send_buffer[seq]
                    if earliest_unacked is None or pkt['sent_time'] < earliest_unacked['sent_time']:
                        earliest_unacked = {'seq': seq, 'sent_time': pkt['sent_time']}

            # 只有最早未确认包超时才重传整个窗口
            if earliest_unacked:
                elapsed = current_time - earliest_unacked['sent_time']
                if elapsed > self.timeout_threshold:
                    print(
                        f"[CLIENT] 检测到超时 (包{earliest_unacked['seq']}等待{elapsed * 1000:.2f}ms > {self.timeout_threshold * 1000:.2f}ms)")
                    print(f"[CLIENT] 重传所有未确认包 (seq={self.base_seq}到{self.next_seq - 1})")
                    for seq in range(self.base_seq, self.next_seq):
                        if seq in self.send_buffer and not self.send_buffer[seq]['acked']:
                            self._retransmit_packet(seq)
                    return True
            return False

    def adjust_timeout(self, new_rtt):
        """RFC 6298标准超时调整"""
        # 第一次测量初始化
        if self.srtt is None:
            self.srtt = new_rtt
            self.rttvar = new_rtt / 2
        else:
            # 计算绝对误差
            err = abs(self.srtt - new_rtt)

            # 更新RTT变化量
            self.rttvar = (1 - self.BETA) * self.rttvar + self.BETA * err

            # 更新平滑RTT
            self.srtt = (1 - self.ALPHA) * self.srtt + self.ALPHA * new_rtt

        # 计算RTO (单位: 毫秒),界定颗粒度防止过小
        rto_ms = self.srtt + max(self.G, self.K * self.rttvar)

        # 设置合理范围 (100ms-3000ms)，防止极端情况
        rto_ms = max(100, min(rto_ms, 3000))

        # 转换为秒并存储
        self.timeout_threshold = rto_ms / 1000.0

        # 调试信息 (减少打印频率)
        if random.random() < 0.1:  # 10%概率打印
            print(f"[RTO更新] SRTT={self.srtt:.2f}ms, "
                  f"RTTVAR={self.rttvar:.2f}ms, "
                  f"RTO={rto_ms:.2f}ms")

    def handle_server_fin(self, seq_num):
        """处理服务器的FIN"""
        print(f"[CLIENT] 收到服务器FIN，seq={seq_num}")

        # 发送ACK
        ack_packet = self.create_packet(
            seq_num=0,
            ack_num=seq_num + 1,
            flags='ACK'
        )
        self.client_socket.sendto(ack_packet, self.server_address)
        print(f"[CLIENT] 发送ACK，ack={seq_num + 1}")

        # 设置结束事件
        self.fin_event.set()

    def close(self):
        """关闭连接"""
        if not self.connected:
            return

        print("[CLIENT] 开始关闭连接...")

        # 发送FIN
        fin_seq = self.next_seq
        fin_packet = self.create_packet(
            seq_num=self.next_seq,
            flags='FIN'
        )
        self.client_socket.sendto(fin_packet, self.server_address)
        print(f"[CLIENT] 发送FIN, seq={self.next_seq}")


        #发送ack
        # 设置接收超时
        self.client_socket.settimeout(2.0)

        try:
            # 第二次挥手：等待服务器的ACK
            data, addr = self.client_socket.recvfrom(1024)
            if addr != self.server_address:
                print(f"[CLIENT] 收到非服务器来源的响应，忽略")
            else:
                seq_num, ack_num, flags, _ = self.parse_packet(data)
                if flags == 'ACK' and ack_num == fin_seq + 1:
                    print(f"[CLIENT] 收到服务器的ACK, ack={ack_num} (第二次挥手)")

                    # 第三次挥手：等待服务器的FIN
                    data, addr = self.client_socket.recvfrom(1024)
                    if addr != self.server_address:
                        print(f"[CLIENT] 收到非服务器来源的响应，忽略")
                    else:
                        seq_num, ack_num, flags, _ = self.parse_packet(data)
                        if flags == 'FIN':
                            print(f"[CLIENT] 收到服务器的FIN, seq={seq_num} (第三次挥手)")

                            # 第四次挥手：发送ACK确认
                            ack_packet = self.create_packet(
                                seq_num=ack_num,
                                ack_num=seq_num + 1,
                                flags='ACK'
                            )
                            self.client_socket.sendto(ack_packet, self.server_address)
                            print(f"[CLIENT] 发送ACK, ack={seq_num + 1} (第四次挥手)")

                            # 等待2MSL时间（2倍最大段生命周期）
                            time.sleep(0.1)  # 简化的MSL等待
        except socket.timeout:
            print("[CLIENT] 关闭连接超时，强制关闭")

        self.client_socket.close()
        self.connected = False
        print("[CLIENT] 连接已关闭")

    def create_packet(self, seq_num, ack_num=0, flags='', data=''):
        """创建数据包，格式：seq|ack|flags|data"""
        return f"{seq_num}|{ack_num}|{flags}|{data}".encode()

    def parse_packet(self, packet):
        """解析数据包"""
        try:
            parts = packet.decode().split('|', 3)
            seq_num = int(parts[0])
            ack_num = int(parts[1])
            flags = parts[2]
            data = parts[3] if len(parts) > 3 else ''
            return seq_num, ack_num, flags, data
        except:
            return 0, 0, '', ''

    def print_summary(self):
        """打印汇总信息"""
        print("\n=== GBN协议传输统计 ===")
        print(f"[CLIENT] - 总发送包数: {self.packets_sent}")
        print(f"[CLIENT] - 重传包数: {self.packets_retransmitted}")

        if self.packets_sent > 0:
            loss_rate = (self.packets_retransmitted / self.packets_sent) * 100
            print(f"[CLIENT] - 丢包率: {loss_rate:.2f}%")
        else:
            print(f"[CLIENT] - 丢包率: 0.00%")

        if self.rtt_samples:
            df = pd.DataFrame(self.rtt_samples, columns=['RTT'])
            print(f"[CLIENT] - 最大RTT: {df['RTT'].max():.2f} ms")
            print(f"[CLIENT] - 最小RTT: {df['RTT'].min():.2f} ms")
            print(f"[CLIENT] - 平均RTT: {df['RTT'].mean():.2f} ms")
            if len(self.rtt_samples) > 1:
                print(f"[CLIENT] - RTT标准差: {df['RTT'].std():.2f} ms")
            else:
                print(f"[CLIENT] - RTT标准差: N/A (样本不足)")
        else:
            print(f"[CLIENT] - 没有收集到RTT样本")

        if self.start_time:
            total_time = time.time() - self.start_time
            print(f"[CLIENT] - 总传输时间: {total_time:.2f}秒")


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python udpclient_gbn.py <server_ip> <server_port>")
        sys.exit(1)
    
    server_ip = sys.argv[1]
    server_port = int(sys.argv[2])
    client = ReliableUDPClient(server_ip, server_port)
    # client = ReliableUDPClient("127.0.0.1", 3083)
    try:
        if client.connect():
            client.start_transfer()
    except KeyboardInterrupt:
        print("\n[CLIENT] 用户中断程序")
        client.close()
    except Exception as e:
        print(f"[CLIENT] 程序错误: {e}")
        import traceback

        traceback.print_exc()
        client.close()